# 搜索雷达水平波束宽度计算
from typing import Dict, List
import json
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from apps.cpm.rlm_context import RlmContext
from apps.cpm.cpm_engine import CpmEngine
from core.radar_core import RadarCore

class SrAziBwAgent(object):
    def __init__(self):
        self.name = 'apps.cpm.agent.sr_azi_bw_agent.SrAziBwAgent'

    @staticmethod
    def run(context:RlmContext) -> Dict:
        
        '''
        作为整体任务第一步时，以用户自然语言为输入，输出为本步完成后的结果。
        '''
        observes = SrAziBwAgent._observe(context=context)
        plans = SrAziBwAgent._reason(observes=observes, context=context)
        return SrAziBwAgent._act(observes=observes, context=context, plans=plans)
    
    @staticmethod
    def _observe(context:RlmContext) -> Dict:
        observes = {}
        prompt = SrAziBwAgent._gen_params_parse_prompt(user_query=context.current_input)
        resp = CpmEngine.infer(query=prompt)
        print(f'resp json: \n{resp[8:-3]};')
        params = json.loads(resp[8:-3])
        context.states['params'].update(params)
        observes.update(params)
        return observes
    
    @staticmethod
    def _reason(observes:Dict, context: RlmContext) -> List:
        plans = []
        return plans
    
    @staticmethod
    def _act(observes:Dict, context: RlmContext, plans: List) -> Dict:
        # 提取参数
        snr = float(observes['snr'])
        tsc = float(observes['tsc'])
        sigma_tgtm = float(observes['sigma_tgtm'])
        sigma_tgta = float(observes['sigma_tgta'])
        range_m = float(observes['range_m'])
        te = float(observes['te'])
        nf = float(observes['nf'])
        loss = float(observes['loss'])
        az_angle = float(observes['az_angle'])
        el_angle = float(observes['el_angle'])
        delta_R = float(observes['delta_R'])
        c = 3.0e8
        # 计算带宽
        tau = 2*delta_R / c
        B = 1 / tau
        # 计算搜索立体角 (球面度)
        omega = (az_angle / 57.296) * (el_angle / 57.296)
        omega_db = 10.0 * np.log10(omega)  # 转换为dB
        # 导弹的PAP计算
        pap_missile = RadarCore.power_aperture(
            snr, tsc, 10**(sigma_tgtm/10), range_m, te, nf, loss, az_angle, el_angle
        )
        # 飞机的PAP计算
        pap_aircraft = RadarCore.power_aperture(
            snr, tsc, 10**(sigma_tgta/10), range_m, te, nf, loss, az_angle, el_angle
        )
        # 计算变化距离下的PAP曲线
        ranges_km = np.arange(2, 91, 0.1)  # 2到90公里，步长1公里
        ranges_m = ranges_km * 1000  # 转换为米
        # 初始化数组
        pap_m = np.zeros_like(ranges_km)
        pap_a = np.zeros_like(ranges_km)
        # 计算每个距离点的PAP
        for i, range_val in enumerate(ranges_m):
            pap_m[i] = RadarCore.power_aperture(
                snr, tsc, 10**(sigma_tgtm/10), range_val, te, nf, loss, az_angle, el_angle
            )
            pap_a[i] = RadarCore.power_aperture(
                snr, tsc, 10**(sigma_tgta/10), range_val, te, nf, loss, az_angle, el_angle
            )
        # 创建恒定PAP值的数组（用于绘制水平参考线）
        missile_pap_const = np.full_like(ranges_km, pap_missile)
        aircraft_pap_const = np.full_like(ranges_km, pap_aircraft)
        # 绘制结果
        plt.figure(figsize=(10, 6))
        # 绘制PAP随距离变化的曲线
        plt.plot(ranges_km, pap_m, 'k', label='Missile')
        plt.plot(ranges_km, pap_a, 'k-.', label='Aircraft')
        # 绘制固定距离处的水平参考线
        plt.plot(ranges_km, missile_pap_const, 'k:', linewidth=1.5)
        plt.plot(ranges_km, aircraft_pap_const, 'k:', linewidth=1.5)
        # 添加图例、标签和网格
        plt.legend()
        plt.xlabel('Range (km)')
        plt.ylabel('Power Aperture Product (dB)')
        plt.title('Power Aperture Product vs Range')
        plt.grid(True)
        # 添加标注
        plt.text(70, pap_missile+1, f'Missile PAP @60km: {pap_missile:.1f} dB', 
                fontsize=10, ha='center')
        plt.text(70, pap_aircraft-1, f'Aircraft PAP @60km: {pap_aircraft:.1f} dB', 
                fontsize=10, ha='center')
        # 显示图形
        plt.tight_layout()
        # plt.show()
        # 计算水平扫描波束宽度（基于天线扫描时间和搜索角度）
        # 波束宽度可以通过扫描时间和搜索角度计算
        # 假设天线扫描速率为：搜索角度/扫描时间
        scan_rate_deg_per_sec = az_angle / tsc  # 度/秒
        scan_rate_rad_per_sec = np.deg2rad(scan_rate_deg_per_sec)  # 弧度/秒
        print(f"\n水平扫描波束宽度相关参数:")
        print(f"扫描角度: {az_angle}°")
        print(f"扫描时间: {tsc}秒")
        print(f"扫描速率: {scan_rate_deg_per_sec:.2f}°/秒 ({scan_rate_rad_per_sec:.4f} rad/s)")
        #
        A_e = 1.75
        PAP = pap_missile
        P_av = 10.0**(PAP/10.0) / A_e
        rho = 0.8
        A = A_e / rho
        f0 = 2.0E9
        G = (4*np.pi*A_e) / ((c / f0)**2)
        G_db = 10*np.log10(G)
        print(f'平均功率为：{P_av}; 孔径：{A}; 增益：{G}; {G_db};')
        theta_e = el_angle
        theta_a = 26000 / (theta_e * G)
        print(f'水平扫描波束宽度：{theta_a}度; {np.deg2rad(theta_a)}; ?')
        # 更新上下文对象
        context.states['params'].update(observes)
        # 生成试验报告
         # 创建目录: work/reports/v001
        base_folder = 'work/reports/e003'
        Path(base_folder).mkdir(parents=True, exist_ok=True)
        # 创建图片目录：work/reports/v001/images
        images_folder = f'{base_folder}/images'
        Path(images_folder).mkdir(parents=True, exist_ok=True)
        plt.savefig(f'{images_folder}/f00.png')
        with open(f'{base_folder}/report.md', 'w', encoding='utf-8') as wfd:
            wfd.write(f'# 雷达设计：确定水平方向波束宽度\n')
            wfd.write(f'# 1. 问题描述\n{context.current_input}\n')
            wfd.write(f'# 2. 计算过程\n')
            wfd.write(f'## 2.1. 计算带宽和脉冲宽度\n')
            wfd.write(f'距离精度为$\\bigtriangleup R=150m$，由$\\bigtriangleup R=\\frac{{c\\tau}}{{2}}=\\frac{{c}}{{2B}} \\Rightarrow \\tau=\\frac{{1}}{{B}}$，可得$\\tau=1\\mu sec$，因此$B=1MHz$\n')
            wfd.write(f'## 2.2. 计算功率孔径积\n')
            wfd.write(f'想要求功率孔径积，必须先计算在探测距离下的最大搜索空间。由需求可得水平扫描范围为$\\Theta_{{A}}={{az_angle}}^{{\\circ}}$，竖直扫描范围为$\\Theta_{{E}}={{el_angle}}^{{\\circ}}$，雷达的搜索空间为：\n')
            wfd.write(f'$$\n\\Omega = \\frac{{\\Theta_{{A}}\\Theta_{{E}}}}{{57.296^{{2}}}}=\\frac{{360 \\times 10}}{{27.296^{{2}}}}=1.097 \\quad \\text{{steradians}}\n$$\n')
            wfd.write(f'飞机功率孔径积：$PAP_{{aircraft}}={pap_aircraft} dB$\n')
            wfd.write(f'导弹功率孔径积：$PAP_{{missile}}={pap_missile} dB$\n')
            wfd.write(f'## 2.3. 求平均发射功率\n')
            wfd.write(f'找出更有挑战性场景，此处为导弹探测，$PAP=PAP_{{missile}}={pap_missile}dB$，我们来求平均发射功率：\n')
            wfd.write(f'$$\nP_{{av}} \\times A_{{e}} = 10^{{\\frac{{PAP}}{{10}}}} = 10^{{3.853}}={P_av}\n$$\n')
            wfd.write(f'## 2.4. 功率孔径积与距离曲线\n')
            wfd.write(f'![f00.png](./images/f00.png "")\n')
            wfd.write(f'## 2.5. 水平扫描波束宽度\n')
            wfd.write(f'$$\nG = k\\frac{{4\\pi}}{{\\theta_{{e}}\\theta_{{a}}}} \\approx \\frac{{26000}}{{\\theta_{{e}}\\theta_{{a}}}}\n$$\n')
            wfd.write(f'$$\n\\theta_{{a}} = \\frac{{26000}}{{\\theta_{{e}}G}}=\\frac{{26000}}{{10 \\times 977.38}}= {theta_a}^{{\\circ}} = {np.deg2rad(theta_a)}\n$$\n')
        return {'result': 0, 'params': {}}
    
    @staticmethod
    def _gen_params_parse_prompt(user_query:str) -> str:
        template = '''
# 角色
假设你是资深雷达系统设计师

# 任务
你需要根据用户输入，提取出所有的参数，并转换为标准单位，以json格式输出。由于参数很多，而且参数表述和单位形式多样，你需要按照下面的步骤执行。

# 要求
你只需要给出json结果，不需要生成代码或进行解释。

# 步骤
1. 信噪比（snr）：名称还可能为SNR，单位为dB（分贝），如果数值不带dB，需要转为dB，公式为：dB=10*np.log10(snr)，提取值为v1；
2. 天线扫描时间（tsc）：名称还可能为扫描率，单位为s（秒）、ms（毫秒）、us（微秒），如果为ms时需乘以e-3，如果为us时需乘以e-6，提取值为v2；
3. 导弹RCS（sigma_tgtm）：名称还可能为导弹截面积、导弹雷达截面积，单位为dBsm，提取值为v3；
4. 飞机RCS（sigma_tgta）：名称还可能为飞机截面积、飞机雷达截面积，单位为dBsm，提取值为v4；
5. 灵敏度距离（range_m）：名称还可能为最大探测距离，单位为m（米）、Km（千米，公里），当单位为Km时需乘以e3，提取值为v5；
6. 有效噪声温度（te）：名称还可能为工作温度，单位为K（开尔文），如果为度或摄氏度需加上275.12，提取值为v6；
7. 噪声系数（nf）：单位为dB（分贝），如果数值不带dB，需要转为dB，公式为：dB=10*np.log10(nf)，提取值为v7；
8. 系统损失（loss）：名称还可能为雷达损失、接收噪声，单位为dB（分贝），如果数值不带dB，需要转为dB，公式为：dB=10*np.log10(loss)，提取值为v8；
9. 水平搜索角度（az_angle）：又叫水平扫描范围，单位为$\\circ$（度），如果只有数值则表示弧度需乘以(180/np.pi)，提取值为v9；
10. 竖直搜索角度（el_angle）：又叫竖直扫描范围，单位为$\\circ$（度），如果只有数值则表示弧度需乘以(180/np.pi)，提取值为v10；
11. 距离精度（delta_R）：单位为m（米）、Km（千米，公里），当单位为Km时需乘以e3，提取值为v11；
11. 生成json结果：
{
    "snr": v1,
    "tsc": v2,
    "sigma_tgtm": v3,
    "sigma_tgta": v4,
    "range_m": v5,
    "te": v6,
    "nf": v7,
    "loss": v8,
    "az_angle": v9,
    "el_angle": v10,
    "delta_R": v11
}

# 示例
输入：设计一个地基雷达，可以识别10公里外的飞机和2公里外的导弹，最大探测距离为60公里。假设飞机的平均RCS为6dBsm，导弹平均RCS为-10dBsm。雷达水平搜索范围为$360^{\\circ}$，竖直搜索范围为$10^{\\circ}$。扫描率为2秒。距离精度为150米。假设噪声数为8dB，接收噪声为10dB，工作温度为290.0K。使用fan beam，水平波束宽度小于3度，SNR为15dB。请确定水平波束宽度。
输出：
{
    "snr": 15,
    "tsc": 2.0,
    "sigma_tgtm": -10.0,
    "sigma_tgta": 6.0,
    "range_m": 60e3,
    "te": 290.0,
    "nf": 8,
    "loss": 10,
    "az_angle": 360.0,
    "el_angle": 10,
    "delta_R": 150
}

# 当前输入
'''
        prompt = f'{template}{user_query}'
        return prompt